/*
    概述：通过三线程搭建流水线计算yolov5s模型检测测试集的mAP
    输出：各类别对应下的txt文件，保存内容为xmin，ymin，xmax和ymax
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <dlfcn.h>
#include <unistd.h>
#include <queue>
#include <vector>
#include <thread>
#include <dirent.h>
#include <mutex>
#include <iostream>
#include <fstream>

#include "opencv2/opencv.hpp"
#include "rknn_api.h"
#include "postprocess.h"

#define PERF_WITH_POST 1

//定义线程间传递数据的队列数据结构和互斥锁
using namespace cv;
using namespace std;
mutex mtxQueueReadImage;
queue<pair<pair<string,Mat>,Mat>> queueReadImage;

mutex mtxQueueOutput;
queue<pair<pair<string,Mat>,vector<rknn_output>>> queueOutput;


double __get_us(struct timeval t) { return (t.tv_sec * 1000000 + t.tv_usec); }

#define Min(a,b) (a<b?a:b)
cv::Mat letterBox(cv::Mat orig_img, int height, int width)
{
    cv::Mat temp, resimg;
    int img_width = orig_img.cols;
    int img_height = orig_img.rows;

    float r = Min((float)height/img_height,(float)width/img_width);
    
    int new_unpad_height = int(img_height * r);
    int new_unpad_width = int(img_width * r);

    int dh = (height - new_unpad_height);
    int dw = (width -new_unpad_width);

    float pad_h = float(dh) / 2;
    float pad_w = float(dw) / 2;

    int top = (int)(pad_h + 0.5);
    int bottom = (int)(pad_h);
    int left = (int)(pad_w + 0.5);
    int right = (int)(pad_w);
    cv::resize(orig_img,temp,cv::Size(new_unpad_width,new_unpad_height), cv::INTER_LINEAR);
    cv::copyMakeBorder(temp, resimg, top, bottom, left, right,cv::BORDER_CONSTANT,(114,114,114));
    // cv::imwrite("makeBorder.jpg",resimg);
    return resimg;
}

static void printRKNNTensor(rknn_tensor_attr *attr)
{
    printf("index=%d name=%s n_dims=%d dims=[%d %d %d %d] n_elems=%d size=%d "
           "fmt=%d type=%d qnt_type=%d fl=%d zp=%d scale=%f\n",
           attr->index, attr->name, attr->n_dims, attr->dims[3], attr->dims[2],
           attr->dims[1], attr->dims[0], attr->n_elems, attr->size, 0, attr->type,
           attr->qnt_type, attr->fl, attr->zp, attr->scale);
}
rknn_tensor_attr output_attrs[3];

/*-------------------------------------------
                  read image
-------------------------------------------*/
void ReadImage(int cpu_id, string image_dir_path)
{   
    
    struct timeval start_time, stop_time;
    gettimeofday(&start_time, NULL);
    if(cpu_id == 1)
    {   
        cpu_set_t mask;
        CPU_ZERO(&mask);
        CPU_SET(cpu_id,&mask);
        if (pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask) < 0)
		    cerr << "set thread affinity failed" << endl;

	    printf("Bind read_image process to CPU %d\n", cpu_id);
    }


	//==================================================================================//
	// 遍历所有需要yolov5s检测的jpg文件
    DIR* dp;
    struct dirent *filename;
    dp = opendir(image_dir_path.c_str());
    if(!dp) cout << "open error!"<<endl;
    vector<string> file_list;

    // cout << "ReadImage test" <<endl;

    while(filename = readdir(dp))
    {
        string s = filename->d_name;
        long id = s.find(".");
        if((id!=0) and (id !=-1))
        {
            if(s.substr(id,4) == ".jpg")
                file_list.push_back(s);
            else continue;
        }
        else continue;
    }
    cout <<"在"<< image_dir_path <<"目录下共找到"<<file_list.size()<<"个jpg文件"<<endl;

    for(int i=0;i < file_list.size();i++)
    {   
        string pic_name = image_dir_path + file_list[i];
        Mat img = imread(pic_name,1);
      
        cv::Mat resimg = letterBox(img, 640,640);

        //cv::cvtColor(resimg, resimg, cv::COLOR_BGR2RGB);
        pair<string,Mat> img_info_pair = make_pair(pic_name,img);
        //将读入图片和rknn_input 放入pair中，然后放入队列，之前先锁住
        mtxQueueReadImage.lock();
        queueReadImage.push(make_pair(img_info_pair,resimg));
        mtxQueueReadImage.unlock();
    }
    gettimeofday(&stop_time, NULL);
    printf(" read-image-and-resize线程所用时间为： %f ms\n",(__get_us(stop_time) - __get_us(start_time)) / 1000);
}


/*-------------------------------------------
                RunRknn
-------------------------------------------*/

void RunRknn(int cpu_id,string task_num, string model_path)
{   
    int ret = 0;
    struct timeval start_time, stop_time;
    gettimeofday(&start_time, NULL);
    if(cpu_id == 2)
    {   
        cpu_set_t mask;
        CPU_ZERO(&mask);
        CPU_SET(cpu_id,&mask);
        if (pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask) < 0)
		    cerr << "set thread affinity failed" << endl;

	    printf("Bind run_rknn process to CPU %d\n", cpu_id);
    }
    
	// Load model
	FILE *fp = fopen(model_path.c_str(), "rb");
    	if(fp == NULL) {
        	printf("fopen %s fail!\n", model_path);
        	return;
    	}
    	fseek(fp, 0, SEEK_END);   //fp指向end,fseek(FILE *stream, long offset, int fromwhere);
    	int model_len = ftell(fp);   //相对文件首偏移
    	void *model = malloc(model_len);
    	fseek(fp, 0, SEEK_SET);   //SEEK_SET为文件头
    	if(model_len != fread(model, 1, model_len, fp)) {
        	printf("fread %s fail!\n", model_path);
        	free(model);
        	return;
    	}
        fclose(fp);
	//init
	rknn_context ctx = 0;
	ret = rknn_init(&ctx,model,model_len,0);
	if(ret < 0) 
    {
        	printf("rknn_init fail! ret=%d\n", ret);
        	return;
    }
    //input
    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = 640 * 640 * 3;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].pass_through = 0;
    //output
    rknn_output outputs[3];
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < 3; i++)
    {
        outputs[i].want_float = 0;                  //需要将输出数据设置为float，否则就不会检测出目标  
    }

    int n = atoi(task_num.c_str());

    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < 3; i++)
    {
        output_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]),
                            sizeof(rknn_tensor_attr));
        printRKNNTensor(&(output_attrs[i]));
    }

    while(true)
    {   
        if(n == 0)
        {
            cout << "rknn process finish!"<<endl;
            break;
        }
        mtxQueueReadImage.lock();
        if(queueReadImage.empty())
        {   
            mtxQueueReadImage.unlock();
            usleep(1000);
        }
        else
        {   
            //获取从read线程传入的队列的数据
            pair<pair<string,Mat>,Mat> image_pair = queueReadImage.front();
            queueReadImage.pop();
            mtxQueueReadImage.unlock();
            inputs[0].buf = image_pair.second.data;
            ret = rknn_inputs_set(ctx, 1, inputs);
            if(ret != 0)
                 cout <<"rknn_inputs_set ret:" << ret <<endl;
            ret = rknn_run(ctx, NULL);
            if(ret != 0)
                 cout <<"rknn_run ret:" << ret <<endl;
            ret = rknn_outputs_get(ctx, 3, outputs, NULL);
            if(ret != 0)
                 cout <<"rknn_outputs_get ret:" << ret <<endl;
            vector<rknn_output> outputVec;
            for(int i =0; i< 3;i++)
            {
                outputVec.push_back(outputs[i]);
            }
            //将rknn出来的数据存放至下一个队列中
            mtxQueueOutput.lock();
            queueOutput.push(make_pair(image_pair.first,outputVec));
            mtxQueueOutput.unlock();
            n--;
	    }
    }
    rknn_outputs_release(ctx, 3, outputs);
    ret = rknn_destroy(ctx);
    gettimeofday(&stop_time, NULL);
    printf(" RunRknn线程所用时间为： %f ms\n",(__get_us(stop_time) - __get_us(start_time)) / 1000);
}
/*-------------------------------------------
                Post_process
-------------------------------------------*/
void PostProcess(int cpu_id, string task_num)
{   
    int width = 640;
    int height =640;
    const float vis_threshold = 0.1;
    const float nms_threshold = 0.5;
    const float conf_threshold = 0.3;

    struct timeval start_time, stop_time;
    gettimeofday(&start_time, NULL);
    if(cpu_id == 3)
    {   
        cpu_set_t mask;
        CPU_ZERO(&mask);
        CPU_SET(cpu_id,&mask);
        if (pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask) < 0)
		    cerr << "set thread affinity failed" << endl;

	    printf("Bind post process to CPU %d\n", cpu_id);
    }
    int n = atoi(task_num.c_str());
    
    //创建*.txt多文件，保存yolov5前向推理的结果，用于计算mAP
    FILE **fps = 0;
	fps = (FILE **)calloc(4, sizeof(FILE *));
    fps[0] = fopen("./print_detector_detections/yawning.txt", "w+");
    fps[1] = fopen("./print_detector_detections/eye_closing.txt", "w+");
    fps[2] = fopen("./print_detector_detections/smoking.txt", "w+");
    fps[3] = fopen("./print_detector_detections/phone_using.txt", "w+");	
    while(true)
    {
        if(n==0)
        {
            cout << "post process finish!"<<endl;
            break;
        }
        rknn_output outputs[3];
        memset(outputs, 0, sizeof(outputs));
        mtxQueueOutput.lock();
        if(queueOutput.empty())
        {
            mtxQueueOutput.unlock();
            usleep(1000);
        }
        else
        {
            pair<pair<string,Mat>,vector<rknn_output>> img_pair;
            vector<rknn_output> outputVec;
            img_pair = queueOutput.front();
            queueOutput.pop();
            mtxQueueOutput.unlock();
            outputVec = img_pair.second;
            pair<string,Mat> img_info_pair = img_pair.first;
            string img_name = img_info_pair.first;
            Mat img = img_info_pair.second;

            int img_width = img.cols;
            int img_height =img.rows;
            for(int i =0;i<3;i++)
            {	
                outputs[i] = outputVec[i];   
            }
            //==================================================================================//
            // YOLOv5 POST Process    
            detect_result_group_t detect_result_group;

            //post process
            float scale_w = (float)width / img_width;
            float scale_h = (float)height / img_height;
            //图片自适应缩放后的h缩放比例，等同于scale_w
            scale_h = scale_w;
            float offset_h =0.5*(640 - img_height * scale_w); 

            vector<float> out_scales;
            vector<uint8_t> out_zps;
            for (int i = 0; i < 3; ++i)
            {
                out_scales.push_back(output_attrs[i].scale);
                out_zps.push_back(output_attrs[i].zp);
            }

            post_process((uint8_t *)outputs[0].buf, (uint8_t *)outputs[1].buf, (uint8_t *)outputs[2].buf, height, width,
                        conf_threshold, nms_threshold, vis_threshold, scale_w, scale_h, out_zps, out_scales, &detect_result_group, offset_h);
            

            // Draw Objects
            int cnt =  detect_result_group.count;
            if(cnt == 0)
                printf("%s没有检测到对象\n",img_name.c_str());
            else    
                printf("%s检测效果如下\n",img_name.c_str());
            for (int i = 0; i < cnt; i++)
            {   
                detect_result_t *det_result = &(detect_result_group.results[i]);
                printf("%s @ (%d %d %d %d) %f\n",
                    det_result->name,
                    det_result->box.left, det_result->box.top, det_result->box.right, det_result->box.bottom,
                    det_result->prop);
                int x1 = det_result->box.left;
                int y1 = det_result->box.top;
                int x2 = det_result->box.right;
                int y2 = det_result->box.bottom;
                // rectangle(img, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(255, 0, 0, 255), 3);
                // putText(img, det_result->name, cv::Point(x1, y1 + 12), 1, 2, cv::Scalar(0, 255, 0, 255));

                //将检测框对应的图片名 置信度以及四个坐标输出到各个类别对应的txt文件中 
                int j = 0;
                if(strcmp(det_result->name, "yawning") == 0)   j = 0;
                else if(strcmp(det_result -> name, "eye_closing") == 0) j = 1;
                else if(strcmp(det_result -> name, "smoking")== 0) j = 2;
                else if(strcmp(det_result -> name, "phone_using") == 0) j =3;
                else;
                string name = img_name.substr(img_name.rfind("/")+1, img_name.rfind(".")-img_name.rfind("/") -1);
				fprintf(fps[j], "%s %f %d %d %d %d\n", name.c_str(), det_result->prop,
                        x1, y1, x2, y2);
            }
            n--;
        }
    }

	for(int j = 0; j < 4; ++j)
		if(fps[j]) fclose(fps[j]);

    gettimeofday(&stop_time, NULL);
    printf("PostProcess 线程所用时间为： %f ms\n",(__get_us(stop_time) - __get_us(start_time)) / 1000);   
}

/*-------------------------------------------
                 Main Function
-------------------------------------------*/
int main(int argc , char*argv[4]) 
{
	if (argc != 4) {
		cout << "Usage of this exec: ./程序名 需要检测的jpg文件夹路径 jpg数量 rknn模型文件" << endl;
		return -1;
	}
	// vector<string> labels = readCocoNames(coco_names);

    struct timeval start_time, stop_time;
    gettimeofday(&start_time, NULL);

    int i, cpus = 0;
	int camera_index;
	cpu_set_t mask;
	cpu_set_t get;
	array<thread, 3> threads;

	cpus = sysconf(_SC_NPROCESSORS_CONF);     //处理器个数是4
	printf("This system has %d processor(s).\n", cpus);

	threads = {thread(ReadImage,1,argv[1]),
               thread(RunRknn,2,argv[2],argv[3]),
			   thread(PostProcess,3,argv[2])};
	cout << "this is main thread" <<endl;
	for (int i = 0; i < 3; i++)
		threads[i].join();

    gettimeofday(&stop_time, NULL);
    printf(" 整个进程的总计时长是： %f ms\n",(__get_us(stop_time) - __get_us(start_time)) / 1000);
}